import os

import cv2

from src.utils import _args_checker


def img2video(img_dir, img_size, video_dir, fps):
    print('正在生成视频......')
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')  # opencv3.0
    videoWriter = cv2.VideoWriter(video_dir, fourcc, fps, img_size)

    f_l=[]
    for idx in os.listdir(img_dir):
        if idx.split(".")[-1] == "png":
            f_l.append(idx)

    for idx in sorted(f_l,key=lambda x: int(x[4:-4])):
        img = os.path.join(img_dir, idx)
        frame = cv2.imread(img)
        videoWriter.write(frame)

    videoWriter.release()
    print('转换完成！')

def clean_pics(dir):
    for idx in sorted(os.listdir(dir)):
        if idx.split(".")[-1] == "png":
            os.remove(os.path.join(dir, idx))


def generate_video(**kwargs):
    save_pictures = _args_checker(kwargs, "save_pictures", allow_type=bool, default=True)
    w = _args_checker(kwargs,"w",allow_type=int,default=1920)
    h = _args_checker(kwargs,"h",allow_type=int,default=1080)
    fps = _args_checker(kwargs, "fps", allow_type=int, default=25)
    path = _args_checker(kwargs, "source_path",allow_type=str, allow_none=False)
    outpath = _args_checker(kwargs, "output_path",allow_type=str, default=os.path.join(path,"output"))
    video_name = _args_checker(kwargs, "video_name",allow_type=str, default="video.avi")

    img_size = (w, h)  # w, h
    img2video(img_dir=outpath, img_size=img_size, video_dir=os.path.join(outpath,video_name), fps=fps)

    if not save_pictures:
        clean_pics(outpath)
